package com.ho1ho.audioexample.utils;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.util.Log;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;

/**
 * Created by Michael Leo <y@ho1ho.com><br>
 * Date: 2020/02/03 14:48
 * <p>
 * Code is far away from bug with the animal protecting
 * </p>
 * <pre>
 * ----------Dragon be here!----------/
 *    ┏┓   ┏┓
 *   ┏┛┻━━━┛┻┓
 *   ┃       ┃
 *   ┃   ━   ┃
 *   ┃ ┳┛ ┗┳ ┃
 *   ┃       ┃
 *   ┃   ┻   ┃
 *   ┃       ┃
 *   ┗━┓   ┏━┛
 *     ┃   ┃神兽保佑
 *     ┃   ┃代码无BUG！
 *     ┃   ┗━━━┓
 *     ┃       ┣┓
 *     ┃       ┏┛
 *     ┗┓┓┏━┳┓┏┛
 *      ┃┫┫ ┃┫┫
 *      ┗┻┛ ┗┻┛
 * ━━━━━━神兽出没━━━━━━
 * </pre>
 */
public class AacEncoder {
    private static final String TAG = AacEncoder.class.getSimpleName();

    private MediaCodec.BufferInfo mBufferInfo;
    private long mPresentationTimeUs = 0;
    private ByteArrayOutputStream mOutputAacStream = new ByteArrayOutputStream();
    private MediaCodec mAudioEncoder;

    public AacEncoder(int sampleRate, int bitrate, int channelCount) {
        try {
            mAudioEncoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_AUDIO_AAC);
        } catch (IOException e) {
            e.printStackTrace();
        }

        MediaFormat mediaFormat = MediaFormat.createAudioFormat(MediaFormat.MIMETYPE_AUDIO_AAC,
                sampleRate, channelCount);

//        mediaFormat.setString(MediaFormat.KEY_MIME, MediaFormat.MIMETYPE_AUDIO_AAC);
        mediaFormat.setInteger(MediaFormat.KEY_AAC_PROFILE, MediaCodecInfo.CodecProfileLevel.AACObjectLC);
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, bitrate);
//        mediaFormat.setInteger(MediaFormat.KEY_CHANNEL_MASK, DEFAULT_AUDIO_FORMAT);
        mediaFormat.setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, 1024 * 1024);
        mAudioEncoder.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
        mAudioEncoder.start();

        mBufferInfo = new MediaCodec.BufferInfo();
    }

    /**
     * Encode pcm data to aac.<br>
     * <br>
     *
     * @param pcmData PCM data.
     */
    public byte[] encodePcmToAac(byte[] pcmData) throws Exception {
        Log.d(TAG, pcmData.length + " is coming");

        // See the dequeueInputBuffer method in document to confirm the timeoutUs parameter.
        int inputBufferIndex = mAudioEncoder.dequeueInputBuffer(-1);
        if (inputBufferIndex >= 0) {
            ByteBuffer inputBuffer = mAudioEncoder.getInputBuffer(inputBufferIndex);
            if (inputBuffer != null) {
                inputBuffer.clear();
                inputBuffer.put(pcmData);
//            inputBuffer.limit(pcmData.length);
            }
            long pts = computePresentationTimeUs(mPresentationTimeUs);
            mAudioEncoder.queueInputBuffer(inputBufferIndex, 0, pcmData.length, pts, 0);
            mPresentationTimeUs += 1;
        }

        int outputBufferIndex = mAudioEncoder.dequeueOutputBuffer(mBufferInfo, 0);
        while (outputBufferIndex >= 0) {
            ByteBuffer outputBuffer = mAudioEncoder.getOutputBuffer(outputBufferIndex);
            if (outputBuffer != null) {
                int outAacDataSize = mBufferInfo.size;
                // The length of ADTS header is 7.
                int outAacDataSizeWithAdts = outAacDataSize + 7;
                outputBuffer.position(mBufferInfo.offset);
                outputBuffer.limit(mBufferInfo.offset + outAacDataSize);

                // Add ADTS header to pcm data array which its length is pcm data length plus 7.
                byte[] outAacDataWithAdts = new byte[outAacDataSizeWithAdts];
                addAdtsToData(outAacDataWithAdts, outAacDataSizeWithAdts);

                outputBuffer.get(outAacDataWithAdts, 7, outAacDataSize);
                outputBuffer.position(mBufferInfo.offset);
                mOutputAacStream.write(outAacDataWithAdts);
            }

//            Log.e("AudioEncoder", outAacDataWithAdts.length + " bytes written: " + Arrays.toString(outAacDataWithAdts));
//            callback.onCallback(outAacDataWithAdts);

            mAudioEncoder.releaseOutputBuffer(outputBufferIndex, false);
            outputBufferIndex = mAudioEncoder.dequeueOutputBuffer(mBufferInfo, 0);
        }

        mOutputAacStream.flush();
        byte[] outAacBytes = mOutputAacStream.toByteArray();
        mOutputAacStream.reset();

        return outAacBytes;
    }

    /**
     * Add 7-bits ADTS header to aac audio data.<br>
     * <br>
     * https://blog.csdn.net/tx3344/article/details/7414543
     * https://blog.csdn.net/jay100500/article/details/52955232
     * https://wiki.multimedia.cx/index.php/MPEG-4_Audio#Sampling_Frequencies
     *
     * <br>
     * profile
     * 1: Main profile
     * 2: Low Complexity profile(LC)
     * 3: Scalable Sampling Rate profile(SSR)
     * <br>
     * sampling_frequency_index
     * 0: 96000 Hz
     * 1: 88200 Hz
     * 2: 64000 Hz
     * 3: 48000 Hz
     * 4: 44100 Hz
     * 5: 32000 Hz
     * 6: 24000 Hz
     * 7: 22050 Hz
     * 8: 16000 Hz
     * 9: 12000 Hz
     * 10: 11025 Hz
     * 11: 8000 Hz
     * 12: 7350 Hz
     * 13: Reserved
     * 14: Reserved
     * 15: frequency is written explictly
     *
     * <br>
     * channel_configuration
     * 0: Defined in AOT Specifc Config
     * 1: 1 channel: front-center
     * 2: 2 channels: front-left, front-right
     * 3: 3 channels: front-center, front-left, front-right
     * 4: 4 channels: front-center, front-left, front-right, back-center
     * 5: 5 channels: front-center, front-left, front-right, back-left, back-right
     * 6: 6 channels: front-center, front-left, front-right, back-left, back-right, LFE-channel
     * 7: 8 channels: front-center, front-left, front-right, side-left, side-right, back-left, back-right, LFE-channel
     * 8-15: Reserved
     *
     * @param outAacDataWithAdts    The audio data with ADTS header.
     * @param outAacDataLenWithAdts The length of audio data with ADTS header.
     */
    private void addAdtsToData(byte[] outAacDataWithAdts, int outAacDataLenWithAdts) {
        int profile = 2; // AAC LC. If you change this value, DO NOT forget to change KEY_AAC_PROFILE while config MediaCodec
        int freqIdx = 8; // 4: 44.1KHz 8: 16Khz 11: 8Khz
        int chanCfg = 1; // 1: single_channel_element 2: CPE(channel_pair_element)
        outAacDataWithAdts[0] = (byte) 0xFF;
        outAacDataWithAdts[1] = (byte) 0xF9;
        outAacDataWithAdts[2] = (byte) (((profile - 1) << 6) + (freqIdx << 2) + (chanCfg >> 2));
        outAacDataWithAdts[3] = (byte) (((chanCfg & 3) << 6) + (outAacDataLenWithAdts >> 11));
        outAacDataWithAdts[4] = (byte) ((outAacDataLenWithAdts & 0x7FF) >> 3);
        outAacDataWithAdts[5] = (byte) (((outAacDataLenWithAdts & 7) << 5) + 0x1F);
        outAacDataWithAdts[6] = (byte) 0xFC;
    }

    /**
     * Release sources.
     */
    public void close() {
        try {
            if (mAudioEncoder != null) {
                mAudioEncoder.stop();
                mAudioEncoder.release();
            }
            if (mOutputAacStream != null) {
                mOutputAacStream.flush();
                mOutputAacStream.close();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * Calculate PTS.<br>
     * <br>
     * Actually, it doesn't make any error if you return 0 directly.
     *
     * @return The calculated presentation time in microseconds.
     */
    public static long computePresentationTimeUs(long frameIndex) {
        return frameIndex * 1_000_000 / 16_000;
    }
}